
#include "BuildinMath.h"

#include "OpenglCore.h"


#ifndef UTILS_SHADER
#define UTILS_SHADER

#include <vector>
#include <fstream>
#include <unordered_map>

enum class ShaderType { Vertex, Fragment, Geometry };

static ShaderType ToShaderType(const std::string& type) {
    if (type == "vertex")
        return ShaderType::Vertex;
    if (type == "fragment" || type == "pixel")
        return ShaderType::Fragment;
    if (type == "geometry" || type == "pixel")
        return ShaderType::Geometry;
    throw std::runtime_error("Invalid shader type");
}

static GLenum ToGLenum(const ShaderType type) {
    switch (type) {
        default: throw std::runtime_error("Invalid shader type"); break;
        case ShaderType::Vertex:    return GL_VERTEX_SHADER;
        case ShaderType::Fragment:  return GL_FRAGMENT_SHADER;
        case ShaderType::Geometry:  return GL_GEOMETRY_SHADER;
    }
}

class Shader {
public:
    uint32_t renderID = 0;
    std::unordered_map<std::string, uint32_t> uniformLocationCache;

    Shader() = default;
    virtual ~Shader() { GLCALL(glDeleteProgram(renderID)); }

	void load(const char* path) {
		std::unordered_map<ShaderType, std::string> shaderSource;
		std::string stringSource;

		std::ifstream in(path, std::ios::in | std::ios::binary);
		if (!in) {
			throw std::runtime_error("Failed to parse shader source");
			return;
		}
		in.seekg(0, std::ios::end);
		stringSource.resize(in.tellg());
		in.seekg(0, std::ios::beg);
		in.read(&stringSource[0], stringSource.size());
		in.close();

		parseShaderSource(shaderSource, stringSource);
		compileShaders(shaderSource);
	}

    void setUniformInt  (const std::string& name, const int value)   { GLCALL(glUniform1i(getUniformLocation(name), value)); }
    void setUniformFloat(const std::string& name, const float value) { GLCALL(glUniform1f(getUniformLocation(name), value)); }
    void setUniformMat3f(const std::string& name, const Matrix3f& m) { GLCALL(glUniformMatrix3fv(getUniformLocation(name), 1, GL_FALSE, &m[0][0])); }
    void setUniformMat4f(const std::string& name, const Matrix4f& m) { GLCALL(glUniformMatrix4fv(getUniformLocation(name), 1, GL_FALSE, &m[0][0])); }
    void setUniformVec3f(const std::string& name, const Vector3f& v) { GLCALL(glUniform3f(getUniformLocation(name), v[0], v[1], v[2])); }
    void setUniformVec4f(const std::string& name, const Vector4f& v) { GLCALL(glUniform4f(getUniformLocation(name), v[0], v[1], v[2], v[3])); }

    void bind() const { GLCALL(glUseProgram(renderID)); }
    void unbind() const { GLCALL(glUseProgram(0)); }

private:
    bool parseShaderSource(std::unordered_map<ShaderType, std::string>& shaderSource, const std::string& stringSource) {
        const char* typeToken = "#type";
		int64_t typeTokenLength = std::strlen(typeToken);
		int64_t pos = stringSource.find(typeToken, 0);

		while (pos != std::string::npos) {
			int64_t eol = stringSource.find_first_of("\r\n", pos);
            if (eol == std::string::npos)
                return false;  // Shader source code format error!

			int64_t begin = pos + typeTokenLength + 1;
			std::string type = stringSource.substr(begin, eol - begin);

			int64_t nextLinePos = stringSource.find_first_not_of("\r\n", eol);
			pos = stringSource.find(typeToken, nextLinePos);

			auto shaderType = ToShaderType(type);
            if (!shaderSource[shaderType].empty())
                return false;  // This shader is already exist!

			shaderSource[shaderType] = stringSource.substr(nextLinePos, pos - (nextLinePos == std::string::npos ? stringSource.size() - 1 : nextLinePos));
		}
        return true;
    }

    bool compileShaders(const std::unordered_map<ShaderType, std::string>& shaderSource) {
        GLCALL(renderID = glCreateProgram());

		int32_t shader = 0;
		for (auto source : shaderSource) {
			const char* src = source.second.c_str();
			GLenum type = ToGLenum(source.first);
			GLCALL(shader = glCreateShader(type));
			GLCALL(glShaderSource(shader, 1, &src, 0));
			GLCALL(glCompileShader(shader));

			int success = 0;
			GLCALL(glGetShaderiv(shader, GL_COMPILE_STATUS, &success));
			if (!success) {
				int maxLength = 0;
				GLCALL(glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &maxLength));

				std::vector<GLchar> infoLog(maxLength);
				GLCALL(glGetShaderInfoLog(shader, maxLength, &maxLength, &infoLog[0]));

				GLCALL(glDeleteShader(shader));
                throw std::runtime_error(infoLog.data());
                return false;
			}
			GLCALL(glAttachShader(renderID, shader));
		}

		GLCALL(glLinkProgram(renderID));

		GLint isLinked = 0;
		glGetProgramiv(renderID, GL_LINK_STATUS, &isLinked);
		if (isLinked == GL_FALSE) {
			GLint maxLength = 0;
			GLCALL(glGetProgramiv(renderID, GL_INFO_LOG_LENGTH, &maxLength));

			std::vector<GLchar> infoLog(maxLength);
			GLCALL(glGetProgramInfoLog(renderID, maxLength, &maxLength, &infoLog[0]));

			GLCALL(glDeleteProgram(renderID));

			for (auto source : shaderSource)
				GLCALL(glDeleteShader(shader));

            throw std::runtime_error(infoLog.data());
            return false;
		}
		for (auto source : shaderSource)
			GLCALL(glDeleteShader(shader));
        return true;
    }

    int getUniformLocation(const std::string& name) {
		if (uniformLocationCache.find(name) != uniformLocationCache.end())
			return uniformLocationCache[name];

		GLCALL(int location = glGetUniformLocation(renderID, name.c_str()));
		if (location == -1)
            std::cout << "uniform '" << name << "' does not exist !" << std::endl;
		uniformLocationCache[name] = location;

		return location;
    }
};

#endif // UTILS_SHADER
